import numpy as np
import pandas as pd
from math import *
p=lambda m,s,x:(1/(sqrt(2*pi)*s))*exp(-(x-m)**2/2*s**2)

def getData():
    '''
        获取数据集
        :return:
        '''
    csv_data = pd.read_csv("watermelon_dataset.csv")
    data = [[c, r, k, t, u, j,d,s,g] for c, r, k, t, u, j, d, s,g in
            zip(csv_data['色泽'], csv_data['根蒂'], csv_data['敲声'], csv_data['纹理'], csv_data['脐部'], csv_data['触感'],
                csv_data['密度'],csv_data['含糖率'],csv_data['好瓜'])]
    label = ['色泽', '根蒂', '敲声', '纹理', '脐部', '触感','密度','含糖率']
    attributeType={}
    for i,l in zip(range(len(label)+1),label):
        attributeType[l]=set([feat[i] for feat in data])
    return data, label,attributeType

def calConditionalProb(data,attributeType):
    # 类先验概率
    classPriorProb={}
    classPriorProb['是']=len([feat[8] for feat in data if feat[8]=='是' ])/len(data)
    classPriorProb['否']=1-classPriorProb['是']
    #print(classPriorProb)

    # 条件概率
    conditionProb={}
    for i,a in zip(range(len(attributeType)-1),attributeType):
        for t in attributeType[a]:
            conditionProb[t] = {}
            conditionProb[t]['是'] = len([feat[8] for feat in data if feat[i] == t and feat[8] == '是'])/len([feat[8] for feat in data if feat[8]=='是' ])
            conditionProb[t]['否'] = len([feat[8] for feat in data if feat[i] == t and feat[8] == '否'])/len([feat[8] for feat in data if feat[8]=='否' ])
    msm=np.array(list(attributeType['密度']))[:8].mean()
    msv=np.array(list(attributeType['密度']))[:8].var()
    mfm = np.array(list(attributeType['密度']))[9:].mean()
    mfv = np.array(list(attributeType['密度']))[9:].var()
    hsm = np.array(list(attributeType['含糖率']))[:8].mean()
    hsv = np.array(list(attributeType['含糖率']))[:8].var()
    hfm = np.array(list(attributeType['含糖率']))[9:].mean()
    hfv = np.array(list(attributeType['含糖率']))[9:].var()


    for i in list(attributeType['密度'])[:8]:
        conditionProb[i]['是']=p(msm,msv,i)
    for i in list(attributeType['密度'])[9:]:
        conditionProb[i]['否']=p(mfm,mfv,i)

    #for i in list(attributeType['含糖率'])[:8]:
    #    conditionProb[i]['是']=p(hsm,hsv,i)
    #for i in list(attributeType['含糖率'])[9:]:
    #   conditionProb[i]['否']=p(hfm,hfv,i)

    print(conditionProb)
    return classPriorProb,conditionProb

if __name__ =='__main__':
    data,label,attributeType=getData()
    print(attributeType)
    calConditionalProb(data,attributeType)
